import copy

import numpy as np
import scipy as sp

import ice
import routs


class ShipSpeed:
    """
    Models the ship speed in the presence of sea ice.

    Attributes:
        ice_cover: Model of the ice cover.
        ice_class: Ice class of the ship.
        mw60: Indicated that the ship is a 60 MW icebreaker.
        speed_no_ice: Maximum speed with no ice at sea (knots).
    """
    ice_cover: ice.IceCover
    ice_class: int
    mw60: bool
    speed_no_ice: float

    def __init__(self, ice_cover: ice.IceCover, ice_class: int, mw60: bool, speed_no_ice: float):
        """
        Args:
            ice_cover: Model of the ice cover.
            ice_class: Ice class of the ship.
            mw60: Indicated that the ship is a 60 MW icebreaker.
            speed_no_ice: Maximum speed with no ice at sea (knots).
        """
        self.ice_cover = ice_cover
        if ice_class in {0, 1, 2, 3, 4, 5, 6, 7, 9}:
            self.ice_class = ice_class
        else:
            raise ValueError("invalid ice class")
        self.mw60 = mw60
        self.speed_no_ice = speed_no_ice

    def __str__(self):
        if self.mw60:
            return f"60 MW icebreaker ({self.speed_no_ice:.1f} knots)"
        elif self.ice_class == 0:
            return f"none ({self.speed_no_ice:.1f} knots)"
        else:
            return f"Arc {self.ice_class} ({self.speed_no_ice:.1f} knots)"

    def compute(self, speed_60mw: float, assisted: bool=True) -> float:
        """
        Compute the maximum ship speed.

        Args:
            speed_60mw: Maximum speed of a 60 MW icebreaker at the point of interest.
            assisted: Indicates that the ship is assisted by an icebreaker.
        Returns:
            The maximum ship speed (knots).
        """

        # Fast return.
        if self.mw60:
            return speed_60mw
        # Compute the speed reduction.
        if 20 <= speed_60mw:
            return self.speed_no_ice
        elif 15 <= speed_60mw:
            if self.ice_class <= 3:
                if assisted:
                    return self.speed_no_ice
                else:
                    return 0.0
            elif self.ice_class <= 6:
                if assisted:
                    return 0.8 * self.speed_no_ice
                else:
                    return 0.0
            elif self.ice_class <= 7:
                if assisted:
                    return 0.6 * self.speed_no_ice
                else:
                    return self.speed_no_ice
            else:
                return 0.9 * speed_60mw
        elif 10 <= speed_60mw:
            if self.ice_class <= 3:
                return 0.0
            elif self.ice_class <= 6:
                if assisted:
                    return 0.7 * self.speed_no_ice
                else:
                    return 0.0
            elif self.ice_class <= 7:
                if assisted:
                    return 0.15 * self.speed_no_ice
                else:
                    return 0.0
            else:
                return 0.75 * speed_60mw
        else:
            return 0.0


class ShipRoute:
    """
    Describes all possible routs of a ship.
    """

    def __init__(self, ship_speed: ShipSpeed, graph: routs.Graph):
        """
        Args:
            ship_speed: Ship speed model.
            graph: Graph.
        """

        # Copy the nodes.
        self.nodes = graph.nodes

        # Select a subset of edges.
        self.edges = []
        for edge in graph.edges:
            # Copy the edge.
            edge_copy = copy.deepcopy(edge)

            # Recompute the speed.
            for i in range(len(edge.speed_60mw)):
                speed_this = ship_speed.compute(edge.speed_60mw[i], assisted=False)
                if speed_this == 0.0:
                    break
            else:
                self.edges.append(edge_copy)

        # Compute the weight matrix.
        weights = np.zeros((len(self.nodes), len(self.nodes)))
        for edge in self.edges:
            weights[edge.node_start.index, edge.node_end.index] = edge.length

        # Find the shortest path for each pair of nodes.
        self.distances, self.predecessors =\
            sp.sparse.csgraph.shortest_path(weights, method="auto", directed=False, return_predecessors=True)

    def find_path(self, index_start: int, index_end: int) -> list[routs.Edge]:
        """
        Find the fastest path from a one node to another.

        Args:
            index_start: Index of the start node.
            index_end: Index of the end node.

        Returns:
            A sequence of edges forming the fastest path. If the end node is inaccessible then
            the sequence is empty.
        """

        # Check that the end point is accessible.
        if np.isinf(self.distances[index_start, index_end].flat[0]):
            return []

        # Reconstruct the path.
        path_indices = [index_end]
        index = index_end
        while index != index_start and index != -9999:
            index = self.predecessors[index_start, index]
            path_indices.append(index)
        path_indices = list(reversed(path_indices))

        # Form the sequence of edges.
        path_edges = []
        for i in range(1, len(path_indices)):
            for edge in self.edges:
                if edge.node_start.index == path_indices[i - 1] and edge.node_end.index == path_indices[i]:
                    path_edges.append(edge)

        return path_edges
